import os
os.environ['JAX_PLATFORMS'] = 'cpu'

import jax
import numpy as np
import jax.numpy as jnp
from problems.jax.cec2022 import CEC2022TestSuit


# 其余代码...
class CEC2022Functions:
    """CEC2022测试函数集的简易接口"""
    
    def __init__(self):
        # 初始化随机密钥
        self.key = jax.random.PRNGKey(0)
        
        # 预先初始化所有12个函数实例并保存其状态
        self.problems = {}
        self.states = {}
        for i in range(1, 13):
            problem = CEC2022TestSuit.create(i)
            # 正确初始化状态
            state = problem.init(self.key)
            self.problems[i] = problem
            self.states[i] = state
            
    def __getitem__(self, func_num):
        """允许使用all_functions[i]形式调用第i个函数"""
        if func_num not in self.problems:
            raise ValueError(f"函数编号必须在1-12之间，收到了{func_num}")
            
        # 返回一个可调用的函数，该函数接受x并返回适应度
        def evaluate_func(x):
            # 转换为jnp数组
            if isinstance(x, np.ndarray):
                if x.ndim == 1:
                    # 一维数组，转为单个样本的二维数组
                    x_jnp = jnp.array(x).reshape(1, -1)
                elif x.ndim == 2:
                    # 二维数组，直接转换
                    x_jnp = jnp.array(x)
                else:
                    raise ValueError("输入数组维度必须是1或2")
            else:
                raise TypeError("输入必须是numpy数组")
                
            # 评估函数
            problem = self.problems[func_num]
            state = self.states[func_num]
            F, _ = problem.evaluate(state, x_jnp)
            
            # 如果输入是一维的，返回单个适应度值
            if x.ndim == 1:
                return np.array(F)[0]
            else:
                return np.array(F)
                
        return evaluate_func

# 创建可全局访问的接口实例
all_functions_2022 = CEC2022Functions()

if __name__ == "__main__":
    # 单个解评估 (一维数组)
    x = np.random.uniform(-100, 100, 10)  # 10维空间中的一个点
    f1 = all_functions_2022[1](x)  # F1_CEC2022的适应度值
    print(f"F1适应度: {f1}")

    # 多个解评估 (二维数组)
    population = np.random.uniform(-100, 100, (50, 10))  # 50个点，每个点10维
    f2_values = all_functions_2022[2](population)  # F2_CEC2022的适应度值
    print(f"F2最小适应度: {np.min(f2_values)}")